#include "aidetectbingxingkernel.h"


// 设备端插值函数 (辅助函数)
__device__ uchar3 interpolateBilinear(uchar3* src, int width, int height, int step, float x, float y) {
    int x0 = static_cast<int>(x);
    int y0 = static_cast<int>(y);
    int x1 = min(x0 + 1, width - 1);
    int y1 = min(y0 + 1, height - 1);

    float dx = x - x0;
    float dy = y - y0;

    uchar3 p00 = *reinterpret_cast<uchar3*>(src + y0 * step + x0 * 3);
    uchar3 p01 = *reinterpret_cast<uchar3*>(src + y0 * step + x1 * 3);
    uchar3 p10 = *reinterpret_cast<uchar3*>(src + y1 * step + x0 * 3);
    uchar3 p11 = *reinterpret_cast<uchar3*>(src + y1 * step + x1 * 3);

    return make_uchar3(
        static_cast<uchar>((1-dx)*(1-dy)*p00.x + dx*(1-dy)*p01.x + (1-dx)*dy*p10.x + dx*dy*p11.x),
        static_cast<uchar>((1-dx)*(1-dy)*p00.y + dx*(1-dy)*p01.y + (1-dx)*dy*p10.y + dx*dy*p11.y),
        static_cast<uchar>((1-dx)*(1-dy)*p00.z + dx*(1-dy)*p01.z + (1-dx)*dy*p10.z + dx*dy*p11.z)
    );
}

// 预处理核函数实现
__global__ void preprocessKernel(uchar3* src, int srcWidth, int srcHeight, int srcStep,
                               float* dst, int dstWidth, int dstHeight) {
   int x = blockIdx.x * blockDim.x + threadIdx.x;
   int y = blockIdx.y * blockDim.y + threadIdx.y;

   if (x < dstWidth && y < dstHeight) {
       float fx = static_cast<float>(x) * srcWidth / dstWidth;
       float fy = static_cast<float>(y) * srcHeight / dstHeight;

       uchar3 interpolated = interpolateBilinear(src, srcWidth, srcHeight, srcStep, fx, fy);

       int dstIdx = (y * dstWidth + x) * 3;
       dst[dstIdx + 0] = interpolated.x / 127.5f - 1.0f;
       dst[dstIdx + 1] = interpolated.y / 127.5f - 1.0f;
       dst[dstIdx + 2] = interpolated.z / 127.5f - 1.0f;
   }
}

// 后处理核函数实现
__global__ void postprocessKernel(const float* output, int numClasses,
                                uchar* mask, int width, int height) {
   int x = blockIdx.x * blockDim.x + threadIdx.x;
   int y = blockIdx.y * blockDim.y + threadIdx.y;

   if (x < width && y < height) {
       int maxIndex = 0;
       float maxProb = -FLT_MAX;

       for (int c = 0; c < numClasses; c++) {
           float prob = output[c * width * height + y * width + x];
           if (prob > maxProb) {
               maxProb = prob;
               maxIndex = c;
           }
       }

       mask[y * width + x] = (maxIndex == 1) ? 255 : 0;
   }
}

__global__ void hello(){
    int i = blockDim.x * blockIdx.x + threadIdx.x;
    printf("hello");
}

int test_cuda(void){
    hello<<<1 , 3>>>();
    cudaDeviceSynchronize();
    printf("done");
}
